import numpy as np
import platgo as pg


class testGM(pg.Problem):

    # 此测试函数仅用于验证梯度法的正确性，初始决策变量分别设置为[[0,0]],[[2,1]],[[-1.2, 1]]
    # 分别运行167,173,164代后的目标值分别为2.5016e-11,2.8349e-11,2.7487e-11。
    def __init__(self):
        self.name = "testGM"
        self.M = 1
        self.D = 2
        lb = [-10000] * self.D
        ub = [10000] * self.D
        self.borders = np.array([lb, ub])
        super().__init__()

    def cal_obj(self, pop: pg.Population) -> None:
        x = pop.decs
        pop.objv = np.array([4 * (x[:, 0] ** 2 - x[:, 1]) ** 2 + 3 * (x[:, 0] - 1) ** 2]).T

    def g_fun(self, pop: pg.Population):
        x = pop.decs
        return np.array([16*x[:, 0]*(x[:, 0] ** 2 - x[:, 1])+6*(x[:, 0]-1), -8*(x[:, 0]**2 - x[:, 1])]).T

    def get_optimal(self) -> np.ndarray:
        return np.array([[0]])
